#! /usr/bin/env python


import rospy
from sensor_msgs.msg import Imu
import matplotlib.pyplot as plt
import numpy as np

import math


last_time = 0

#https://blog.csdn.net/weixin_42782150/article/details/107015617
def data_callback(msg):
    # rospy.loginfo("received")
    # msg.angular_velocity
    # msg.linear_acceleration

    global last_time

    acc_x.append(msg.linear_acceleration.x)        # 添加 i 的平方到 y 轴的数据中
    acc_y.append(msg.linear_acceleration.y)        # 添加 i 的平方到 y 轴的数据中
    acc_z.append(msg.linear_acceleration.z)        # 添加 i 的平方到 y 轴的数据中
    acc_x.pop(0)
    acc_y.pop(0)
    acc_z.pop(0)

    plt.clf()              # 清除之前画的图
    plt.subplot(2, 1, 1)
    plt.plot(axes_x, acc_x)        # 画出当前 ax 列表和 ay 列表中的值的图形
    plt.plot(axes_x, acc_y)        # 画出当前 ax 列表和 ay 列表中的值的图形
    plt.plot(axes_x, acc_z)        # 画出当前 ax 列表和 ay 列表中的值的图形
    plt.legend(['x: '+str(msg.linear_acceleration.x), 'y: '+str(msg.linear_acceleration.y), 'z: '+str(msg.linear_acceleration.z)], loc='upper left')
    plt.title("Linear Acceleration")

    gyr_x.append(msg.angular_velocity.x)        # 添加 i 的平方到 y 轴的数据中
    gyr_y.append(msg.angular_velocity.y)        # 添加 i 的平方到 y 轴的数据中
    gyr_z.append(msg.angular_velocity.z)        # 添加 i 的平方到 y 轴的数据中
    gyr_x.pop(0)
    gyr_y.pop(0)
    gyr_z.pop(0)

    plt.subplot(2, 1, 2)
    plt.plot(axes_x, gyr_x)        # 画出当前 ax 列表和 ay 列表中的值的图形
    plt.plot(axes_x, gyr_y)        # 画出当前 ax 列表和 ay 列表中的值的图形
    plt.plot(axes_x, gyr_z)        # 画出当前 ax 列表和 ay 列表中的值的图形
    plt.legend(['x: '+str(msg.angular_velocity.x), 'y: '+str(msg.angular_velocity.y), 'z: '+str(msg.angular_velocity.z)], loc='upper left')
    plt.title("Angular Velocity")

    if math.isnan(msg.angular_velocity.x) or math.isnan(msg.angular_velocity.y) or math.isnan(msg.angular_velocity.z):
        rospy.logerr("angular_velocitimety Nan!!!")
        print("angular_velocity Nan!!!")
    if math.isnan(msg.linear_acceleration.x) or math.isnan(msg.linear_acceleration.y)or math.isnan(msg.linear_acceleration.z):
        rospy.logerr("linear_acceleration Nan!!!")
        print("linear_acceleration Nan!!!")

    if last_time == msg.header.stamp.to_sec():
        print("\033[0;31;40m", "time same!!!", "\033[0m")
    last_time = msg.header.stamp.to_sec()
    plt.xlabel(f"time stamp: {last_time}", fontsize=30)


    # plt.pause(0.0001)         # 暂停一秒
    plt.ioff()             # 关闭画图的窗口


if __name__ == "__main__":
    rospy.init_node("Display")
    # rospy.loginfo("Hello World!!!!")

    axes_x = list(range(100))          # 定义一个 x 轴的空列表用来接收动态的数据

    acc_x = list(np.zeros(100))       # 定义一个 y 轴的空列表用来接收动态的数据
    acc_y = list(np.zeros(100))       # 定义一个 y 轴的空列表用来接收动态的数据
    acc_z = list(np.zeros(100))       # 定义一个 y 轴的空列表用来接收动态的数据
    gyr_x = list(np.zeros(100))       # 定义一个 y 轴的空列表用来接收动态的数据
    gyr_y = list(np.zeros(100))       # 定义一个 y 轴的空列表用来接收动态的数据
    gyr_z = list(np.zeros(100))       # 定义一个 y 轴的空列表用来接收动态的数据
    plt.ion()                  # 开启一个画图的窗口

    sub = rospy.Subscriber("/imu/raw_data", Imu, data_callback, queue_size=1)



    rospy.spin()